#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""检测结果处理器

用于优化和增强检测结果，提供置信度过滤、非极大值抑制、目标关联和告警规则应用等功能
"""

import logging
import numpy as np
from typing import List, Dict, Optional, Any, Tuple, Callable
from collections import defaultdict
import time

# 设置日志记录器
logger = logging.getLogger(__name__)

class ResultProcessor:
    """检测结果处理器
    
    提供各种检测结果优化和后处理功能
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """初始化检测结果处理器
        
        Args:
            config: 处理器配置字典
        """
        # 默认配置
        self.config = {
            'confidence_threshold': 0.5,  # 置信度阈值
            'nms_threshold': 0.45,        # NMS阈值
            'min_bbox_area': 0,           # 最小边界框面积
            'max_bbox_area': float('inf'), # 最大边界框面积
            'alert_rules': {},            # 告警规则配置
            'target_classes': None,       # 目标类别过滤（None表示不过滤）
            'bbox_scale_factor': 1.0,     # 边界框缩放因子
            'enable_class_specific_thresholds': False,  # 是否启用类别特定阈值
            'class_specific_thresholds': {},  # 类别特定阈值
            'enable_alert_grouping': True,    # 是否启用告警分组
            'max_alert_group_size': 10,      # 最大告警分组大小
        }
        
        # 更新配置
        if config:
            self.config.update(config)
        
        # 告警历史记录，用于去重
        self.alert_history = defaultdict(list)
        self.alert_history_size = 1000  # 最大告警历史记录数
        self.alert_timeout = 5.0        # 告警去重时间窗口（秒）
        
        # 缓存上次告警时间
        self.last_alert_time = time.time()
        
        logger.info("ResultProcessor初始化完成")
    
    def filter_by_confidence(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据置信度过滤检测结果
        
        Args:
            detections: 检测结果列表
            
        Returns:
            List[Dict]: 过滤后的检测结果
        """
        filtered_detections = []
        
        for detection in detections:
            # 获取置信度
            confidence = detection.get('confidence', 0)
            
            # 检查是否启用类别特定阈值
            if self.config['enable_class_specific_thresholds']:
                class_id = detection.get('class_id', None)
                if class_id is not None and class_id in self.config['class_specific_thresholds']:
                    threshold = self.config['class_specific_thresholds'][class_id]
                else:
                    threshold = self.config['confidence_threshold']
            else:
                threshold = self.config['confidence_threshold']
            
            # 应用过滤
            if confidence >= threshold:
                filtered_detections.append(detection)
        
        return filtered_detections
    
    def nms(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """应用非极大值抑制
        
        Args:
            detections: 检测结果列表
            
        Returns:
            List[Dict]: NMS处理后的检测结果
        """
        if not detections:
            return []
        
        # 按类别分组处理
        class_detections = defaultdict(list)
        for detection in detections:
            class_id = detection.get('class_id', 0)
            class_detections[class_id].append(detection)
        
        nms_detections = []
        
        # 对每个类别分别应用NMS
        for class_id, class_detections_list in class_detections.items():
            # 提取边界框和置信度
            bboxes = np.array([d['bbox'] for d in class_detections_list])
            scores = np.array([d['confidence'] for d in class_detections_list])
            
            # 执行NMS
            indices = self._non_max_suppression(
                bboxes, 
                scores, 
                self.config['nms_threshold']
            )
            
            # 保留选定的检测结果
            for i in indices:
                nms_detections.append(class_detections_list[i])
        
        return nms_detections
    
    def _non_max_suppression(self, bboxes: np.ndarray, scores: np.ndarray, 
                            threshold: float) -> List[int]:
        """非极大值抑制实现
        
        Args:
            bboxes: 边界框数组 [N, 4], 格式: [x1, y1, x2, y2]
            scores: 置信度数组 [N]
            threshold: IoU阈值
            
        Returns:
            List[int]: 保留的索引列表
        """
        if len(bboxes) == 0:
            return []
        
        # 获取边界框坐标
        x1 = bboxes[:, 0]
        y1 = bboxes[:, 1]
        x2 = bboxes[:, 2]
        y2 = bboxes[:, 3]
        
        # 计算面积
        area = (x2 - x1 + 1) * (y2 - y1 + 1)
        
        # 根据置信度排序
        order = scores.argsort()[::-1]
        
        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            
            # 计算与其他框的IoU
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            
            # 计算交集面积
            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            
            # 计算IoU
            iou = inter / (area[i] + area[order[1:]] - inter)
            
            # 保留IoU小于阈值的边界框
            inds = np.where(iou <= threshold)[0]
            order = order[inds + 1]
        
        return keep
    
    def filter_by_area(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据边界框面积过滤检测结果
        
        Args:
            detections: 检测结果列表
            
        Returns:
            List[Dict]: 过滤后的检测结果
        """
        filtered_detections = []
        
        for detection in detections:
            bbox = detection.get('bbox', [0, 0, 0, 0])
            area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
            
            # 检查面积是否在范围内
            if self.config['min_bbox_area'] <= area <= self.config['max_bbox_area']:
                filtered_detections.append(detection)
        
        return filtered_detections
    
    def filter_by_class(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据类别过滤检测结果
        
        Args:
            detections: 检测结果列表
            
        Returns:
            List[Dict]: 过滤后的检测结果
        """
        if self.config['target_classes'] is None:
            return detections
        
        filtered_detections = []
        for detection in detections:
            class_id = detection.get('class_id', None)
            class_name = detection.get('class_name', '').lower()
            
            # 检查类别ID或名称是否匹配
            if (class_id is not None and class_id in self.config['target_classes']) or \
               (class_name in map(str.lower, self.config['target_classes'])):
                filtered_detections.append(detection)
        
        return filtered_detections
    
    def scale_bboxes(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """缩放边界框
        
        Args:
            detections: 检测结果列表
            
        Returns:
            List[Dict]: 缩放后的检测结果
        """
        if self.config['bbox_scale_factor'] == 1.0:
            return detections
        
        scaled_detections = []
        for detection in detections:
            bbox = detection.get('bbox', [0, 0, 0, 0])
            x1, y1, x2, y2 = bbox
            
            # 计算中心点和尺寸
            center_x = (x1 + x2) / 2
            center_y = (y1 + y2) / 2
            width = x2 - x1
            height = y2 - y1
            
            # 缩放尺寸
            new_width = width * self.config['bbox_scale_factor']
            new_height = height * self.config['bbox_scale_factor']
            
            # 计算新的边界框坐标
            new_x1 = max(0, center_x - new_width / 2)
            new_y1 = max(0, center_y - new_height / 2)
            new_x2 = center_x + new_width / 2
            new_y2 = center_y + new_height / 2
            
            # 更新检测结果
            scaled_detection = detection.copy()
            scaled_detection['bbox'] = [int(new_x1), int(new_y1), int(new_x2), int(new_y2)]
            scaled_detections.append(scaled_detection)
        
        return scaled_detections
    
    def apply_alert_rules(self, detections: List[Dict[str, Any]], 
                         image_info: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """应用告警规则
        
        Args:
            detections: 检测结果列表
            image_info: 图像相关信息
            
        Returns:
            List[Dict]: 带有告警信息的检测结果
        """
        detections_with_alerts = []
        alert_groups = defaultdict(list)
        
        for detection in detections:
            # 复制原始检测结果
            detection_with_alerts = detection.copy()
            detection_with_alerts['alerts'] = []
            
            # 应用告警规则
            alerts = self._generate_alerts(detection, image_info)
            
            # 去重告警
            unique_alerts = self._deduplicate_alerts(alerts, detection)
            detection_with_alerts['alerts'] = unique_alerts
            
            # 分组告警
            if self.config['enable_alert_grouping']:
                for alert in unique_alerts:
                    alert_type = alert.get('type', 'default')
                    alert_groups[alert_type].append({
                        'detection': detection_with_alerts,
                        'alert': alert
                    })
            
            detections_with_alerts.append(detection_with_alerts)
        
        # 处理告警分组
        if self.config['enable_alert_grouping']:
            self._process_alert_groups(alert_groups, detections_with_alerts)
        
        return detections_with_alerts
    
    def _generate_alerts(self, detection: Dict[str, Any], 
                        image_info: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据检测结果生成告警
        
        Args:
            detection: 检测结果
            image_info: 图像相关信息
            
        Returns:
            List[Dict]: 生成的告警列表
        """
        alerts = []
        class_id = detection.get('class_id', None)
        class_name = detection.get('class_name', '').lower()
        confidence = detection.get('confidence', 0)
        
        # 遍历告警规则
        for rule_name, rule_config in self.config['alert_rules'].items():
            # 检查规则是否适用
            if self._is_rule_applicable(rule_config, class_id, class_name):
                # 检查规则条件
                if self._check_rule_conditions(rule_config, detection, image_info):
                    # 生成告警
                    alert = {
                        'type': rule_name,
                        'severity': rule_config.get('severity', 'warning'),
                        'message': self._generate_alert_message(rule_config, detection),
                        'timestamp': time.time(),
                        'confidence': confidence,
                        'class_id': class_id,
                        'class_name': class_name,
                        'bbox': detection.get('bbox')
                    }
                    
                    # 添加额外信息
                    if 'extra_info' in rule_config:
                        alert.update(rule_config['extra_info'])
                    
                    alerts.append(alert)
        
        # 添加默认告警规则（可根据需要启用）
        # 默认高置信度告警
        if confidence > 0.9:
            alerts.append({
                'type': 'high_confidence',
                'severity': 'info',
                'message': f"高置信度检测: {class_name} ({confidence:.2f})".strip(),
                'timestamp': time.time(),
                'confidence': confidence,
                'class_id': class_id,
                'class_name': class_name,
                'bbox': detection.get('bbox')
            })
        
        return alerts
    
    def _is_rule_applicable(self, rule_config: Dict[str, Any], 
                           class_id: Optional[int], class_name: str) -> bool:
        """检查规则是否适用于当前检测结果
        
        Args:
            rule_config: 规则配置
            class_id: 类别ID
            class_name: 类别名称
            
        Returns:
            bool: 是否适用
        """
        # 检查适用的类别
        applicable_classes = rule_config.get('classes', [])
        if not applicable_classes:
            return True  # 没有指定类别，默认适用
        
        # 检查类别ID或名称
        if class_id is not None and class_id in applicable_classes:
            return True
        
        if class_name in map(str.lower, applicable_classes):
            return True
        
        return False
    
    def _check_rule_conditions(self, rule_config: Dict[str, Any], 
                              detection: Dict[str, Any], 
                              image_info: Optional[Dict[str, Any]]) -> bool:
        """检查规则条件是否满足
        
        Args:
            rule_config: 规则配置
            detection: 检测结果
            image_info: 图像相关信息
            
        Returns:
            bool: 条件是否满足
        """
        conditions = rule_config.get('conditions', {})
        
        # 检查置信度条件
        if 'min_confidence' in conditions:
            if detection.get('confidence', 0) < conditions['min_confidence']:
                return False
        
        if 'max_confidence' in conditions:
            if detection.get('confidence', 0) > conditions['max_confidence']:
                return False
        
        # 检查边界框条件
        if 'bbox_area' in conditions:
            bbox = detection.get('bbox', [0, 0, 0, 0])
            area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
            area_range = conditions['bbox_area']
            if area < area_range[0] or area > area_range[1]:
                return False
        
        # 这里可以添加更多条件检查
        # 例如：位置条件、时间条件、关联对象条件等
        
        return True
    
    def _generate_alert_message(self, rule_config: Dict[str, Any], 
                               detection: Dict[str, Any]) -> str:
        """生成告警消息
        
        Args:
            rule_config: 规则配置
            detection: 检测结果
            
        Returns:
            str: 告警消息
        """
        # 使用模板或默认格式
        template = rule_config.get('message_template', 
                                 "检测到 {class_name} ({confidence:.2f} 置信度)")
        
        # 替换模板变量
        message = template.format(
            class_name=detection.get('class_name', '未知对象'),
            confidence=detection.get('confidence', 0),
            class_id=detection.get('class_id', '未知'),
            timestamp=time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()),
            bbox=detection.get('bbox', []),
            **detection
        )
        
        return message
    
    def _deduplicate_alerts(self, alerts: List[Dict[str, Any]], 
                           detection: Dict[str, Any]) -> List[Dict[str, Any]]:
        """去重告警
        
        Args:
            alerts: 告警列表
            detection: 检测结果
            
        Returns:
            List[Dict]: 去重后的告警列表
        """
        unique_alerts = []
        current_time = time.time()
        bbox = detection.get('bbox', [0, 0, 0, 0])
        class_id = detection.get('class_id', None)
        
        for alert in alerts:
            alert_type = alert.get('type', 'default')
            
            # 清理过期的告警历史
            self._cleanup_alert_history(current_time)
            
            # 检查是否是重复告警
            is_duplicate = False
            history_key = f"{alert_type}_{class_id}"
            
            for history_alert in self.alert_history[history_key]:
                # 计算边界框距离
                distance = self._calculate_bbox_distance(bbox, history_alert['bbox'])
                time_diff = current_time - history_alert['timestamp']
                
                # 如果在时间窗口内且距离较近，则视为重复告警
                if time_diff < self.alert_timeout and distance < 50:  # 50像素阈值
                    is_duplicate = True
                    break
            
            if not is_duplicate:
                unique_alerts.append(alert)
                # 记录到告警历史
                self.alert_history[history_key].append(alert)
        
        return unique_alerts
    
    def _cleanup_alert_history(self, current_time: float):
        """清理过期的告警历史
        
        Args:
            current_time: 当前时间
        """
        # 限制历史记录大小
        for key in self.alert_history:
            # 删除过期告警
            self.alert_history[key] = [
                alert for alert in self.alert_history[key] 
                if current_time - alert['timestamp'] < self.alert_timeout
            ]
            
            # 限制每个键的记录数量
            if len(self.alert_history[key]) > 100:
                self.alert_history[key] = self.alert_history[key][-100:]
        
        # 定期检查总历史记录数量
        total_size = sum(len(alerts) for alerts in self.alert_history.values())
        if total_size > self.alert_history_size:
            # 删除最旧的历史记录
            oldest_key = None
            oldest_time = current_time
            
            for key, alerts in self.alert_history.items():
                if alerts and alerts[0]['timestamp'] < oldest_time:
                    oldest_time = alerts[0]['timestamp']
                    oldest_key = key
            
            if oldest_key:
                # 删除一部分历史记录
                self.alert_history[oldest_key] = self.alert_history[oldest_key][len(self.alert_history[oldest_key])//2:]
    
    def _calculate_bbox_distance(self, bbox1: List[int], bbox2: List[int]) -> float:
        """计算两个边界框的距离
        
        Args:
            bbox1: 第一个边界框 [x1, y1, x2, y2]
            bbox2: 第二个边界框 [x1, y1, x2, y2]
            
        Returns:
            float: 中心点之间的欧几里得距离
        """
        # 计算中心点
        center1 = [(bbox1[0] + bbox1[2])/2, (bbox1[1] + bbox1[3])/2]
        center2 = [(bbox2[0] + bbox2[2])/2, (bbox2[1] + bbox2[3])/2]
        
        # 计算欧几里得距离
        distance = np.sqrt(
            (center1[0] - center2[0])**2 + 
            (center1[1] - center2[1])** 2
        )
        
        return distance
    
    def _process_alert_groups(self, alert_groups: Dict[str, List[Dict[str, Any]]],
                            detections: List[Dict[str, Any]]):
        """处理告警分组
        
        Args:
            alert_groups: 告警分组字典
            detections: 检测结果列表
        """
        for alert_type, group_items in alert_groups.items():
            # 检查分组大小
            if len(group_items) > self.config['max_alert_group_size']:
                # 对相同类型的告警进行合并或优先级排序
                # 这里简单按置信度排序，只保留置信度最高的几个
                group_items.sort(key=lambda x: x['detection'].get('confidence', 0), reverse=True)
                
                # 记录合并信息
                logger.info(f"合并 {alert_type} 类型的告警，从 {len(group_items)} 个合并到 {self.config['max_alert_group_size']} 个")
    
    def process(self, detections: List[Dict[str, Any]], 
                image_info: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """完整处理流程
        
        Args:
            detections: 原始检测结果
            image_info: 图像相关信息
            
        Returns:
            List[Dict]: 处理后的检测结果
        """
        processed_detections = detections
        
        # 1. 类别过滤
        processed_detections = self.filter_by_class(processed_detections)
        
        # 2. 面积过滤
        processed_detections = self.filter_by_area(processed_detections)
        
        # 3. 置信度过滤
        processed_detections = self.filter_by_confidence(processed_detections)
        
        # 4. 非极大值抑制
        processed_detections = self.nms(processed_detections)
        
        # 5. 边界框缩放
        processed_detections = self.scale_bboxes(processed_detections)
        
        # 6. 应用告警规则
        processed_detections = self.apply_alert_rules(processed_detections, image_info)
        
        return processed_detections
    
    def update_config(self, config: Dict[str, Any]):
        """更新配置
        
        Args:
            config: 新的配置
        """
        self.config.update(config)
        logger.info(f"ResultProcessor配置已更新: {config}")
    
    def get_config(self) -> Dict[str, Any]:
        """获取当前配置
        
        Returns:
            Dict: 当前配置
        """
        return self.config.copy()

# 创建全局结果处理器实例
global_result_processor = ResultProcessor()

def get_result_processor() -> ResultProcessor:
    """获取全局结果处理器实例
    
    Returns:
        ResultProcessor: 全局结果处理器实例
    """
    return global_result_processor